{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/DeepLabCut/DeepLabCut/blob/master/examples/COLAB/COLAB_transformer_reID.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TGChzLdc-lUJ"
   },
   "source": [
    "# DeepLabCut 2.2 Toolbox Demo on how to use our Pose Transformer for unsupervised identity tracking of animals\n",
    "![alt text](https://images.squarespace-cdn.com/content/v1/57f6d51c9f74566f55ecf271/1628250004229-KVYD7JJVHYEFDJ32L9VJ/DLClogo2021.jpg?format=1000w)\n",
    "\n",
    "https://github.com/DeepLabCut/DeepLabCut\n",
    "\n",
    "### This notebook illustrates how to use the transformer for a multi-animal DeepLabCut (maDLC) Demo 3 mouse project:\n",
    "- load our mini-demo data that includes a pretrained model and unlabeled video.\n",
    "- analyze a novel video.\n",
    "- use the transformer to do unsupervised ID tracking.\n",
    "- create quality check plots and video.\n",
    "\n",
    "### To create a full maDLC pipeline please see our full docs: https://deeplabcut.github.io/DeepLabCut/README.html\n",
    "- Of interest is a full how-to for maDLC: https://deeplabcut.github.io/DeepLabCut/docs/maDLC_UserGuide.html\n",
    "- a quick guide to maDLC: https://deeplabcut.github.io/DeepLabCut/docs/tutorial.html\n",
    "- a demo COLAB for how to use maDLC on your own data: https://github.com/DeepLabCut/DeepLabCut/blob/main/examples/COLAB/COLAB_maDLC_TrainNetwork_VideoAnalysis.ipynb\n",
    "\n",
    "### To get started, please go to \"Runtime\" ->\"change runtime type\"->select \"Python3\", and then select \"GPU\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HoNN2_0Z9rr_"
   },
   "outputs": [],
   "source": [
    "# Install the latest DeepLabCut version:\n",
    "!apt update && apt install cuda-11-8\n",
    "!pip install \"deeplabcut[tf]\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Wid0GTGMAEnZ"
   },
   "source": [
    "No information needs edited in the cells below, you can simply click run on each:\n",
    "\n",
    "### Download our Demo Project from our server:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PusLdqbqJi60"
   },
   "outputs": [],
   "source": [
    "# Download our demo project:\n",
    "import requests\n",
    "from io import BytesIO\n",
    "from zipfile import ZipFile\n",
    "\n",
    "url_record = 'https://zenodo.org/api/records/7883589'\n",
    "response = requests.get(url_record)\n",
    "if response.status_code == 200:\n",
    "    file = response.json()['files'][0]\n",
    "    title = file['key']\n",
    "    print(f\"Downloading {title}...\")\n",
    "    with requests.get(file['links']['self'], stream=True) as r:\n",
    "        with ZipFile(BytesIO(r.content)) as zf:\n",
    "            zf.extractall(path='/content')\n",
    "else:\n",
    "    raise ValueError(f'The URL {url_record} could not be reached.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8iXtySnQB0BE"
   },
   "source": [
    "## Analyze a novel 3 mouse video with our maDLC DLCRNet, pretrained on 3 mice data \n",
    "\n",
    "###in one step, since auto_track=True you extract detections and association costs, create tracklets, & stitch them. We can use this to compare to the transformer-guided tracking below.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "odYrU3o8BSAr"
   },
   "outputs": [],
   "source": [
    "import deeplabcut as dlc\n",
    "import os\n",
    "\n",
    "project_path = \"/content/demo-me-2021-07-14\"\n",
    "config_path = os.path.join(project_path, \"config.yaml\")\n",
    "video = os.path.join(project_path, \"videos\", \"videocompressed1.mp4\")\n",
    "\n",
    "dlc.analyze_videos(config_path,[video], shuffle=0, videotype=\"mp4\",auto_track=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zmdSLRTOER00"
   },
   "source": [
    "### Next, you compute the local, spatio-temporal grouping and track body part assemblies frame-by-frame:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "F-d6kXqnGeUP"
   },
   "source": [
    "## Create a pretty video output:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "aTRbuUQ1FBO0"
   },
   "outputs": [],
   "source": [
    "#Filter the predictions to remove small jitter, if desired:\n",
    "dlc.filterpredictions(config_path,\n",
    "                                 [video],\n",
    "                                 shuffle=0,\n",
    "                                 videotype='mp4',\n",
    "                                 )\n",
    "\n",
    "dlc.create_labeled_video(\n",
    "    config_path,\n",
    "    [video],\n",
    "    videotype='mp4',\n",
    "    shuffle=0,\n",
    "    color_by=\"individual\",\n",
    "    keypoints_only=False,\n",
    "    draw_skeleton=True,\n",
    "    filtered=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AYNlrgeNUG4U"
   },
   "source": [
    "Now, on the left panel if you click the folder icon, you will see the project folder \"demo-me..\"; click on this and go into \"videos\" and you can find the \"..._id_labeled.mp4\" video, which you can double-click on to download and inspect!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "n7GWMBJUA9x5"
   },
   "source": [
    "### Create Plots of your data:\n",
    "\n",
    "> after running, you can look in \"videos\", \"plot-poses\" to check out the trajectories! (sometimes you need to click the folder refresh icon to see it). Within the folder, for example, see plotmus1.png to vide the bodyparts over time vs. pixel position.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7w9BDIA7BB_i"
   },
   "outputs": [],
   "source": [
    "dlc.plot_trajectories(config_path, [video], shuffle=0,videotype='mp4')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "l7BJQq7nxHVz"
   },
   "source": [
    "# Transformer for reID\n",
    "\n",
    "while the tracking here is very good without using the transformer, we want to demo the workflow for you! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5xlO6TVYxQWc"
   },
   "outputs": [],
   "source": [
    "dlc.transformer_reID(config_path, [video],\n",
    "                     shuffle=0, videotype='mp4',\n",
    "                     track_method='ellipse',n_triplets=100\n",
    "                     )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uO_yoqN7xiBT"
   },
   "source": [
    "now we can make another video with the transformer-guided tracking:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "MBMbRFEMxmi4"
   },
   "outputs": [],
   "source": [
    "dlc.plot_trajectories(config_path, [video],\n",
    "                      shuffle=0,videotype='mp4',\n",
    "                      track_method=\"transformer\"\n",
    "                      )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vx3e-r1CoXaX"
   },
   "outputs": [],
   "source": [
    "dlc.create_labeled_video(\n",
    "    config_path,\n",
    "    [video],\n",
    "    videotype='mp4',\n",
    "    shuffle=0,\n",
    "    color_by=\"individual\",\n",
    "    keypoints_only=False,\n",
    "    draw_skeleton=True,\n",
    "    track_method=\"transformer\"\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "COLAB_transformer_reID.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
